package com.benshou.bcss.common.trace;

import com.benshou.bcss.common.j17common.annotation.Nullable;
import com.benshou.bcss.common.uuid.UUIDs;
import org.slf4j.MDC;

/**
 * TraceId 上下文信息
 *
 * @author <a href="mailto:liuzhenbang@benshou.com">Liu Zhenbang</a>
 */
public class TraceIdContext {

    /**
     * TraceId 在 MDC 中的 Key
     */
    private static final String MDC_KEY = "traceId";

    /**
     * 设置 TraceId
     *
     * @param traceId TraceId
     */
    public static void set(String traceId) {
        MDC.put(MDC_KEY, traceId);
    }

    /**
     * @return TraceId
     */
    @Nullable
    public static String get() {
        return MDC.get(MDC_KEY);
    }

    /**
     * 移除 TraceId 值
     */
    public static void remove() {
        MDC.remove(MDC_KEY);
    }

    /**
     * 获取 TraceId 值。如无，初始化（并设置）后返回。
     *
     * @return TraceId
     */
    public static String getOrInitialize() {
        var traceId = MDC.get(MDC_KEY);
        if (traceId == null) {
            traceId = generate();
            MDC.put(MDC_KEY, traceId);
        }

        return traceId;
    }

    /**
     * 生成 TraceId 值
     *
     * @return TraceId
     */
    private static String generate() {
        return UUIDs.short15();
    }
}
